# coding=utf8
from __future__ import division
import os
import torch
import torch.utils.data as data
import numpy as np
import pandas as pd
import cv2
from matplotlib import pyplot as plt

idx2class_map = {0: 'defect'}
class2idx_map = {'defect': 0}


class GLSdata(data.Dataset):
    def __init__(self, anno_pd, transforms):
        self.anno_pd = anno_pd
        self.transforms = transforms

        self.img_paths = anno_pd['img_pathes']
        self.patch_pos =  anno_pd['patch_pos']
        self.bboxs_pos = anno_pd['bboxs_pos']
        self.class_map = {'defect': 0}

    def __len__(self):
        return len(self.img_paths)

    def __getitem__(self, item):
        img_path = self.img_paths[item]
        img = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB)

        img = cv2.resize(img, (300, 300))

        patch_pos = self.patch_pos[item]
        bbxs = self.bboxs_pos[item]

        if bbxs == []:
            bbxs = np.empty((0,5))
        bbxs = np.array(bbxs)  #(n, 5) label last

        if self.transforms:
            img, bbxs = self.transforms(img, bbxs)

        # plt.imshow(img)
        # plt.show()
        return torch.from_numpy(img).permute(2,0,1).float(), bbxs, patch_pos


def collate_fn(batch):
    imgs = []
    targets = []
    patch_pos = []

    for sample in batch:
        imgs.append(sample[0])
        targets.append(torch.FloatTensor(sample[1]))
        patch_pos.append(sample[2])

    patch_pos = np.array(patch_pos)
    return torch.stack(imgs, 0), targets, patch_pos

def collate_fn2(batch):
    imgs = []
    targets = []
    patch_pos = []

    for sample in batch:
        imgs.append(sample[0])
        targets.append(sample[1])
        patch_pos.append(sample[2])

    patch_pos = np.array(patch_pos)
    return torch.stack(imgs, 0), targets, patch_pos



if __name__ == '__main__':
    import cPickle
    anno_save_path = '/media/hszc/data1/glass_data/patches/mask_no_crayon/annos.pkl'
    anno = cPickle.load(open(anno_save_path,'rb'))
    GLS = GLSdata(anno_pd=anno, transforms=None)

    a, b, c = GLS[2]
    # print a.size()
    print b
    # print c

    # print [box for box in GLS.bboxs_pos[2] if box[-1]==0][0]